#include <stdio.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netdb.h>
#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>
#include <getopt.h>
#include <arpa/inet.h>
#include <sys/time.h>
#include <sys/resource.h>

#define PROGRAM_NAME "conn_tester"
#define VERSION "0.01"
#define AUTHOR "Ivan Totev <ivan@totev.org>"

#define DEFAULT_CONNECTIONS 50
#define CONNECTION_TIMEOUT 1
#define GET "GET %s HTTP/1.0\nHost: %s\n\n"
#define RAND_GET "GET %s%s HTTP/1.0\nHost: %s\n\n"
#define POST "POST %s HTTP/1.0\nHost: %s\n\n"
#define RAND_POST "POST %s%s HTTP/1.0\nHost: %s\n\n"

struct hostent *h;
int sock;
struct sockaddr_in remote;

int prepare_conn(char *host, char *ip);
int sendreq(char *get);
char *random_string(void);
void usage(void);
int set_limits(unsigned int connections);

int main(int argc, char **argv)
{
	char *get = NULL;
	char *host = NULL;
	char *ip = NULL;
	char request[1024];

	int i, opt;
	int connections = DEFAULT_CONNECTIONS;
	unsigned int rand = 0;
	unsigned int post = 0;
	unsigned int verbose = 0;

	//getopt
	while ((opt = getopt(argc, argv, "h:g:c:i:prvV")) != -1)
	{
		switch (opt)
		{
			case 'h':
				host = strdup(optarg);
			break;
			case 'g':
				get = strdup(optarg);
			break;
			case 'c':
				connections = atoi(optarg);
			break;
			case 'i':
				ip = strdup(optarg);
			break;
			case 'p':
				post = 1;
			break;
			case 'r':
				rand = 1;
			break;
			case 'v':
				verbose++;
			break;
			case 'V':
				printf("%s, version %s\n", PROGRAM_NAME, VERSION);
				return 0;
			default:
				usage();
		}
	}

	if (optind < argc)
		host = strdup(argv[optind]);

	//Check arguments
	if(!host)
		usage();
	if(!get)
		get = strdup("/");

	if (set_limits(connections) == -1)
	{
		fprintf(stderr, "Can't change the limits\nPlease use root user or reduce the number of connections\n");
		return 1;
	}

	int pid[connections];

	//resolve, prepare structures
	if ( prepare_conn(host, ip) != 0)
		return 1;

	printf("%d connections to %s with ip %s\n", connections, host, inet_ntoa(remote.sin_addr));
	//fork processes
 	for(i=0; i<connections;i++)
	{
		pid[i] = fork();
		if (pid[i] == 0)
		{
			if(post)
				snprintf(request, sizeof(request), POST, get, host);
			else	
				snprintf(request, sizeof(request), GET, get, host);

			while(1)
			{
				if (rand)
				{

					snprintf(request, sizeof(request), RAND_GET, get, random_string(), host);
					if (post)
						snprintf(request, sizeof(request), RAND_POST, get, random_string(), host);
				}
	
				if ( sendreq(request) == 0 )
				{
					if(verbose == 1)
						write(1,".",1);
					else if (verbose > 1)
					{
						printf("=========================================\n");
						printf("Send request to\n%s\n", request);
						printf("=========================================\n");
					}

				}
				else
				{
					if(verbose == 1)
						write(1, "f", 1);
					else if(verbose > 1)
					{
						printf("=========================================\n");
						printf("Failed to connect to %s\n%s\n",host, request); 
						printf("=========================================\n");
					}
				}
				fflush(stdout);
			}
			return 0;
		}
	}

	for(i=0; i<connections;i++)
		waitpid(pid[i],0,0);
	
	return 0;
}

int set_limits(unsigned int connections)
{
	struct rlimit rlim;
	if ( getrlimit(RLIMIT_NOFILE, &rlim) == -1)
	{
		perror(PROGRAM_NAME);
		return -1;
	}

	if (rlim.rlim_cur < connections)
	{
		rlim.rlim_cur = connections;
		rlim.rlim_max = connections;
		if (setrlimit(RLIMIT_NOFILE, &rlim) == -1)
		{
			perror(PROGRAM_NAME);
			return -1;
		}
	}

	if ( getrlimit(RLIMIT_NPROC, &rlim) == -1)
	{
		perror(PROGRAM_NAME);
		return -1;
	}

	if (rlim.rlim_cur < connections)
	{
		rlim.rlim_cur = connections;
		rlim.rlim_max = connections;
		if( setrlimit(RLIMIT_NOFILE, &rlim) == -1)
		{
			perror(PROGRAM_NAME);
			return -1;
		}
	}
	return 0;
}

void usage(void)
{
	fprintf(stderr, "Usage: %s [-h hostname] [-i use ip different ip] [-g request] [-c connection] [-r add random suffix] [-p use post] [-v verbose] [-V print version]\n", PROGRAM_NAME);
	exit(1);
}

char *random_string()
{
	int i;
	char alphanum[] = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
	int len = sizeof(alphanum);
	char s[(len-1)];
	for (i=0; i<len; i++)
		s[i] = alphanum[rand() % (len - 1)];
	return strdup(s);
}

int prepare_conn(char *host, char *ip)
{
	memset(&remote, 0, sizeof(remote));
	remote.sin_family = AF_INET;
	remote.sin_port = htons(80);

	if (ip)
	{
		remote.sin_addr.s_addr = inet_addr(ip);
		if (remote.sin_addr.s_addr == INADDR_NONE)
		{
			herror(PROGRAM_NAME);
			return 1;
		}
		return 0;
	}
	h = gethostbyname(host);
	if(!h)
	{
		herror(PROGRAM_NAME);
		return 1;
	}

	memcpy(&remote.sin_addr, h->h_addr, h->h_length);

	return 0;
}

int sendreq(char *get)
{
	struct timeval timeout;
	timeout.tv_sec = CONNECTION_TIMEOUT;
	timeout.tv_usec = 0;

	sock = socket(AF_INET, SOCK_STREAM, 0);
	if(sock == -1)
	{
		perror("create socket");
		return 1;
	}

	if (setsockopt (sock, SOL_SOCKET, SO_RCVTIMEO, (char *)&timeout, sizeof(timeout)) < 0)
	{
		perror("setsockopt failed");
		return 1;
	}
	
	if (setsockopt (sock, SOL_SOCKET, SO_SNDTIMEO, (char *)&timeout, sizeof(timeout)) < 0)
	{
		perror("setsockopt failed");
		return 1;
	}

	if (connect(sock, (struct sockaddr *)&remote, sizeof(remote)) == -1 )
	{
		//perror("connect");
		return 2;
	}
	write(sock, get, strlen(get));
	close(sock);
	return 0;
}
